package com.idormy.sms.forwarder.utils

import android.content.Context
import android.graphics.Bitmap
import android.graphics.PixelFormat
import android.hardware.display.DisplayManager
import android.hardware.display.VirtualDisplay
import android.media.Image
import android.media.ImageReader
import android.media.projection.MediaProjection
import android.media.projection.MediaProjectionManager
import android.os.Handler
import android.os.HandlerThread
import android.view.WindowManager
import java.io.File
import java.io.FileOutputStream
import kotlin.coroutines.resume
import kotlin.coroutines.suspendCoroutine

class ScreenshotUtils(private val context: Context) {

    companion object {
        private const val TAG = "ScreenshotUtils"
        var mediaProjection: MediaProjection? = null
    }

    private var imageReader: ImageReader? = null
    private var virtualDisplay: VirtualDisplay? = null
    private var backgroundThread: HandlerThread? = null
    private var backgroundHandler: Handler? = null

    private fun startBackgroundThread() {
        backgroundThread = HandlerThread("ScreenshotBackground").also { it.start() }
        backgroundHandler = Handler(backgroundThread!!.looper)
    }

    private fun stopBackgroundThread() {
        backgroundThread?.quitSafely()
        try {
            backgroundThread?.join()
            backgroundThread = null
            backgroundHandler = null
        } catch (e: InterruptedException) {
            Log.e(TAG, "Error stopping background thread: ${e.message}")
        }
    }

    suspend fun takeScreenshot(): File? = suspendCoroutine { continuation ->
        if (mediaProjection == null) {
            Log.e(TAG, "MediaProjection not initialized. Screenshot requires user permission.")
            continuation.resume(null)
            return@suspendCoroutine
        }

        try {
            startBackgroundThread()

            val windowManager = context.getSystemService(Context.WINDOW_SERVICE) as WindowManager
            val display = windowManager.defaultDisplay
            val width = display.width
            val height = display.height
            val density = context.resources.displayMetrics.densityDpi

            imageReader = ImageReader.newInstance(width, height, PixelFormat.RGBA_8888, 1)

            val outputFile = File(context.cacheDir, "screenshot_${System.currentTimeMillis()}.jpg")

            virtualDisplay = mediaProjection?.createVirtualDisplay(
                "ScreenCapture",
                width, height, density,
                DisplayManager.VIRTUAL_DISPLAY_FLAG_AUTO_MIRROR,
                imageReader?.surface,
                null,
                backgroundHandler
            )

            // 等待一小段时间让屏幕内容渲染到ImageReader
            backgroundHandler?.postDelayed({
                try {
                    val image = imageReader?.acquireLatestImage()
                    if (image != null) {
                        val bitmap = imageToBitmap(image)
                        image.close()

                        FileOutputStream(outputFile).use { output ->
                            bitmap.compress(Bitmap.CompressFormat.JPEG, 90, output)
                        }
                        bitmap.recycle()

                        Log.d(TAG, "Screenshot saved: ${outputFile.absolutePath}")
                        cleanup()
                        continuation.resume(outputFile)
                    } else {
                        Log.e(TAG, "Failed to acquire image")
                        cleanup()
                        continuation.resume(null)
                    }
                } catch (e: Exception) {
                    Log.e(TAG, "Error saving screenshot: ${e.message}")
                    cleanup()
                    continuation.resume(null)
                }
            }, 100)

        } catch (e: Exception) {
            Log.e(TAG, "Error taking screenshot: ${e.message}")
            cleanup()
            continuation.resume(null)
        }
    }

    private fun imageToBitmap(image: Image): Bitmap {
        val planes = image.planes
        val buffer = planes[0].buffer
        val pixelStride = planes[0].pixelStride
        val rowStride = planes[0].rowStride
        val rowPadding = rowStride - pixelStride * image.width

        val bitmap = Bitmap.createBitmap(
            image.width + rowPadding / pixelStride,
            image.height,
            Bitmap.Config.ARGB_8888
        )
        bitmap.copyPixelsFromBuffer(buffer)

        return if (rowPadding == 0) {
            bitmap
        } else {
            // 裁剪掉padding部分
            val croppedBitmap = Bitmap.createBitmap(bitmap, 0, 0, image.width, image.height)
            bitmap.recycle()
            croppedBitmap
        }
    }

    private fun cleanup() {
        virtualDisplay?.release()
        virtualDisplay = null
        imageReader?.close()
        imageReader = null
        stopBackgroundThread()
    }

    fun release() {
        cleanup()
        mediaProjection?.stop()
        mediaProjection = null
    }
}
